# -*- coding: utf-8 -*-
"""
Created on Mon Feb  1 10:06:25 2021

@author: 59567
"""
import os
import sys
from pandas import read_pickle, DataFrame, read_excel, date_range, concat

sys.path.append('../')
import numpy as np
from tjdutils.utils import check_path, real_path, TjdDate, load_pickle, current_time
from tjd_slicer.tjd_slicer import Slicer
from copy import deepcopy
import joblib
import torch
from constructor.rnns_pytorch import RNNs
from constructor.tcn import TCN
from constructor import darnn

from torch import tensor
from pandas import to_datetime


# from corrector import Corrector

def check_zero(data):
    data = data.replace(to_replace=0, method='ffill')
    return data


def move_model_to_montior(data, by='num', num=0):
    collector_dict = data['dict']
    if by == 'num':
        keys = list(collector_dict.keys())
        args_tp, data_path = collector_dict[keys[num]]
        name = keys[num]
    else:
        args_tp, data_path = collector_dict[by]
        name = by
    # args_tp是元组记录参数用, data_path是一个dataframe的路径,需要用到check_path
    rolling_pd = read_pickle(check_path(data_path))
    return rolling_pd, name


def load_model(s_dict):
    md_path = s_dict['model_path']
    if s_dict['learner'] == 'machine_learning':
        md = joblib.load(check_path(md_path)).model
    elif s_dict['model'] in ('RNN', 'GRU', 'LSTM', 'TCN'):
        md = torch.load(check_path(md_path))
    return md


def use_model_predict(s_dict, md, x_test):
    if s_dict['learner'] == 'machine_learning':
        out = md.predict(x_test[0])
    else:
        if s_dict['model'] in ('RNN', 'GRU', 'LSTM'):
            model = RNNs(model_name=s_dict['model'],
                         input_size=int(s_dict['n_features']),
                         hidden_size=10,
                         batch_first=True,
                         time_step=int(s_dict['x_time_step']), output_size=1)
        elif s_dict['model'] == 'TCN':
            model = TCN(input_size=s_dict['n_features'], output_size=1, num_channels=[30] * 8,
                        kernel_size=2, dropout=0)
        model.load_state_dict(md)
        x_tensor_test = tensor(x_test, dtype=torch.float32)
        out = model(x_tensor_test)
        out = out.detach().numpy()
    return out


class Monitor():
    rolling_pd: DataFrame
    cum_x_yoy_monthly: DataFrame
    y_yoy: DataFrame

    def __init__(self, rolling_pd, cum_x_yoy, y_yoy, name):
        super(Monitor, self).__init__()
        rolling_pd.index = rolling_pd['end_datetime']
        y_yoy.index = to_datetime(y_yoy.index)
        self.rolling_pd = rolling_pd  # 按照时间滚动下来的dataframe
        self.rolling_t = rolling_pd.index.to_list()  # 滚动时间
        self.x_daily = cum_x_yoy  # 日度特征同比
        # self.x_daily.index = [str(s)[:10] for s in self.x_daily.index]
        self.y = y_yoy
        self.name = name

    def replace_pd(self, now_date, end_datetime, selected_features):
        temp_x = deepcopy(self.x_daily.loc[:, selected_features])
        single_temp_li = [temp_x.loc[:, [s]].iloc[:, [0]] for s in selected_features]
        temp_x = concat(single_temp_li, axis=1)
        # 防止修改原始值,deepcopy下
        for feature in selected_features:  # 对每一列dataframe（特征维度）遍历
            # print('from',feature, temp_x.loc[end_datetime, feature])
            # print('now_date', now_date)
            temp_x.loc[end_datetime, feature] = self.x_daily.loc[now_date, feature]
            # 最后一天日期是当天值,实现每天不同
            # print('to', temp_x.loc[end_datetime, feature])
            # print('temp_x.shape', temp_x.shape)
        return temp_x

    def get_date_by_freq(self, end_datetime, freq, retro=False):
        freq = "QM"
        if freq == 'M':
            if end_datetime == self.rolling_t[-1]:
                re = date_range(start=TjdDate(end_datetime).dt['mfd'],
                                end=self.x_daily.index[-1],
                                freq='D')
            else:
                re = date_range(start=TjdDate(end_datetime).dt['mfd'],
                                end=end_datetime,
                                freq='D')
        elif freq == 'Q':
            if retro:
                re = date_range(end=end_datetime, periods=60, freq='M')
            else:
                re = date_range(end=end_datetime, periods=3, freq='M')
        else:
            st = TjdDate(end_datetime).dt['lmfd']
            stst = TjdDate(st).dt['lmfd']
            if end_datetime == self.rolling_t[-1]:
                re = date_range(start=stst, end=self.x_daily.index[-1], freq='D')
            else:
                re = date_range(start=stst,end=end_datetime, freq='D')
        print('enddatetime', end_datetime)
        return [str(t)[:10] for t in re]

    def get_md_and_tensor(self, retro=False, quick=False):
        all_tensors, all_models, all_targets, all_preds = {}, {}, {}, {}
        if quick:
            print("quick is true")
            rolling_t = self.rolling_t[-1:]
        else:
            rolling_t = deepcopy(self.rolling_t)
        for end_datetime in rolling_t:
            end_datetime_tensors = {}
            end_datetime_targets = {}
            s_dict = self.rolling_pd.loc[end_datetime, :].to_dict()
            freq = s_dict['freq']
            # print('end_datetime', end_datetime)
            if end_datetime == self.rolling_t[0]:
                if retro:
                    all_subtime_t = self.get_date_by_freq(end_datetime, freq, True)
                else:
                    all_subtime_t = self.get_date_by_freq(end_datetime, freq)
            else:
                all_subtime_t = self.get_date_by_freq(end_datetime, freq)

            if s_dict['model'] == 'DARNN':
                train_result = read_pickle(s_dict['model_path'])
                train_result_dict = {item["trainer_key"]: item for item in train_result}
                Encoder = train_result_dict[s_dict["key"]]["Encoder"]
                Decoder = train_result_dict[s_dict["key"]]["Decoder"]
                for t in all_subtime_t:
                    temp_x = self.replace_pd(t, end_datetime, list(s_dict['features']))
                    dataset = Slicer(x=temp_x,
                                     y=self.y,
                                     freq=freq,
                                     end_datetime=end_datetime,
                                     periods=int(s_dict['periods']),
                                     test_len=int(s_dict['test_len']),
                                     features=list(s_dict['features']),
                                     n_features=int(s_dict['n_features']),
                                     y_name=s_dict['y_name'])
                    x_test = dataset.xx[
                             -s_dict['x_time_step'] - s_dict['test_len'] - s_dict['y_time_step'] + 2:].to_numpy()
                    x_test = np.array(np.row_stack((x_test, np.zeros((1, x_test.shape[1])))))
                    Y = np.array(
                        dataset.yy[-s_dict['x_time_step'] - s_dict['test_len'] - s_dict['y_time_step'] + 1:].to_numpy())
                    y_p = darnn.test(test_len=s_dict['test_len'], T=s_dict['x_time_step'] + 1, X=x_test,
                                     batch_size=s_dict['batch_size'],
                                     y=Y, device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'),
                                     Encoder=Encoder, Decoder=Decoder)
                    all_preds[t] = y_p[0]
                    end_datetime_targets[t] = y_p
            else:
                md = load_model(s_dict)
                # print("lloading md")
                # print(end_datetime)
                for t in all_subtime_t:
                    # 通过test数据集的日期月底那一天替换成每天/月的,
                    # 实现每天/月调用slicer切出当天/月的tensor
                    temp_x = self.replace_pd(t, end_datetime, list(s_dict['features']))
                    slicer = Slicer(x=temp_x,
                                    y=self.y,
                                    freq=freq,
                                    end_datetime=end_datetime,
                                    periods=int(s_dict['periods']),
                                    test_len=int(s_dict['test_len']),
                                    features=list(s_dict['features']),
                                    n_features=int(s_dict['n_features']),
                                    y_name=s_dict['y_name'])
                    x_test = slicer.x_test
                    end_datetime_tensors[t] = x_test

                    y_p = use_model_predict(s_dict, md, x_test)
                    # print(end_datetime, t, y_p)
                    all_preds[t] = y_p[0]
                    end_datetime_targets[t] = y_p
                all_tensors[end_datetime] = end_datetime_tensors
                all_targets[end_datetime] = end_datetime_targets
                all_models[end_datetime] = md
        return DataFrame(all_preds, index=[self.name]).T, all_tensors


def run_monitor(cum_x_yoy, y_yoy, c_path, num_monitor=1):
    data = load_pickle(check_path(c_path))
    rolling_metrics = data['df']
    rolling_metrics['key'] = rolling_metrics.index
    metrics = {'rolling_acc': False, 'triangle_area': True, 'test_mse': True, 'test_vbp': False}
    re = []
    for k, v in metrics.items():
        print('sorting metrics >>>', k)
        rolling_metrics.sort_values(by=k, ascending=v, inplace=True)
        for name in list(rolling_metrics['key'])[:num_monitor]:
            rolling_pd, name = move_model_to_montior(data, by=name)
            m = Monitor(rolling_pd, cum_x_yoy, y_yoy, name)
            all_preds = m.get_md_and_tensor()
            re.append(all_preds[0])
    all_monitors = concat(re, axis=1)
    data['y_name'] = data['y_name'].replace(":", "~")
    data['y_name'] = data['y_name'].replace("：", "--")
    pth = "../output/monitor/"
    pth = check_path(pth) + data['y_name'] + '/'
    if not os.path.exists(pth):
        os.makedirs(pth)
    pth = pth + data['y_name'] + current_time() + '.xlsx'
    all_monitors.to_excel(check_path(pth))
    return all_monitors


if __name__ == "__main__":
    x_path = "high_freq_all_data_2021-03-10.pkl"
    y_path = "y集合.xlsx"
    c_path = "D:/高频中台/2021_03_31/output/collector_test/collector_test2021_04_01_14_27_01_05.pkl"
    x_cum_yoy_path = real_path(x_path)  # 日频特征
    y_yoy = read_excel(real_path(y_path), index_col=0)  # y集合
    cum_x_yoy = read_pickle(real_path(x_cum_yoy_path))

    cum_x_yoy = check_zero(cum_x_yoy)

    re = run_monitor(cum_x_yoy, y_yoy, c_path)
